import glob
import math
import os
import sys
from fileutil import FileUtil
from gnuplot import GnuPlot
from polyevaluator import PolyEvaluator
from rmse import RMSE
from trainer import Trainer

# Utility class to perform cross validation.
class CrossValidate():

  def __init__(self, input_file_prefix, full_train_file, test_file, true_file):
    # Setup data and files.
    self.file_list = glob.glob(input_file_prefix + '*')
    self.lamdas = []
    self.lamdas.append(-35)
    self.lamdas.extend(range(-30, 1))
    self.true_file = true_file
    self.full_train_file = full_train_file
    self.test_file = test_file

    # Setup utils.
    self.rmse = RMSE()
    self.fileutil = FileUtil()

  def run_validation(self, output_filepath_prefix):
    trainer = Trainer()
    num_files = len(self.file_list)

    # Lambda* for which the average training error is the least.
    lambda_star = 0
    min_test_rmse = sys.maxint 

    # File to record the average train and test rmse
    avg_train_error_filepath = '/tmp/avg_train_rmse.dat'
    avg_test_error_filepath = '/tmp/avg_test_rmse.dat'
    avg_train_error_file = open(avg_train_error_filepath, 'w')
    avg_test_error_file = open(avg_test_error_filepath, 'w')

    # Read values from the test file.
    test_xvalues, test_yvalues = self.fileutil.read([self.test_file])
    
    # Iterate for each lambdas
    for l in self.lamdas:
      train_rmse = 0
      test_rmse = 0

      # Cross validate by skipping one file from the set of files passed.
      for idx in range(0, num_files):
        tmp_files = []
        tmp_files[:] = self.file_list[:]
        tmp_files.pop(idx)
      
        weights = trainer.train(tmp_files, l) 

        xvalues, yvalues = self.fileutil.read(tmp_files)
        xvalues_t, yvalues_t = self.fileutil.read([self.file_list[idx]])
        cur_train_rmse = self.rmse.rmse(weights, xvalues, yvalues)
        cur_test_rmse = self.rmse.rmse(weights, xvalues_t, yvalues_t)

        train_rmse += cur_train_rmse
        test_rmse += cur_test_rmse

      # Compute the average train and test rmse and record them in the file.
      avg_train_rmse = float(train_rmse) / num_files
      avg_test_rmse = float(test_rmse) / num_files
      avg_train_error_file.write('%d %f\n' % (l, avg_train_rmse))
      avg_test_error_file.write('%d %f\n' % (l, avg_test_rmse))

      if avg_test_rmse < min_test_rmse:
        min_test_rmse = avg_test_rmse
        lambda_star = l 

    # Close the average rmse file
    avg_train_error_file.close()
    avg_test_error_file.close()

    # Plot lambda Vs average errors.
    gplot = GnuPlot()
    plot_data_series = {}
    plot_data_series['Avg Train RMSE'] = avg_train_error_filepath
    plot_data_series['Avg Test RMSE'] = avg_test_error_filepath
    title = 'RMSE'
    gplot.plot(plot_data_series, output_filepath_prefix + "_avg_rmse.png", title)

    # Compute the weights for the best lambda - lambda_star
    weight_star = trainer.train(self.file_list, lambda_star) 

    # Plot the training data, true data and learning polynomial
    evaluator = PolyEvaluator(weight_star)
    learned_poly_filepath = '/tmp/learned.data'
    evaluator.write_data(learned_poly_filepath)

    plot_data_series = {}
    plot_data_series['True Data'] = self.true_file
    plot_data_series['Train Data'] = self.full_train_file
    plot_data_series['Learned Polynomial'] = learned_poly_filepath

    title = 'Lambda* = exp(%f)' % lambda_star
    gplot.plot(plot_data_series, output_filepath_prefix + "_lambda_star.png", title)

    train_xvalues, train_yvalues = self.fileutil.read([self.full_train_file])
    train_rmse_star = self.rmse.rmse(weight_star, train_xvalues, train_yvalues)
    test_rmse_star = self.rmse.rmse(weight_star, test_xvalues, test_yvalues)

    return (lambda_star, train_rmse_star, test_rmse_star, weight_star)
